#!/usr/bin/env python3
# ui/video_server.py
# 视频服务器和面板：接收/显示视频流、YOLO标注
USE_YOLO = False

import socket
import threading
import time
import cv2
import numpy as np
from PyQt5.QtWidgets import (
    QWidget,
    QVBoxLayout,
    QHBoxLayout,
    QLabel,
    QComboBox,
    QCheckBox,
)
from PyQt5.QtGui import QPixmap, QImage, QPainter, QFont, QColor
from PyQt5.QtCore import Qt, QTimer
import os

if USE_YOLO:
    from ultralytics import YOLO


import logging

logger = logging.getLogger(__name__)


class VideoServer:
    """视频服务器：监听0.0.0.0:6003，接收多个客户端视频流"""

    def __init__(self):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind(("0.0.0.0", 6003))
        self.server_socket.listen(5)
        self.clients = {}  # {addr: (frame, timestamp)}
        self.running = True
        self.accept_thread = threading.Thread(
            target=self.accept_connections, daemon=True
        )
        self.accept_thread.start()
        logger.info("VideoServer started on 0.0.0.0:6003")

    def accept_connections(self):
        while self.running:
            try:
                client_socket, addr = self.server_socket.accept()
                client_thread = threading.Thread(
                    target=self.handle_client, args=(client_socket, addr), daemon=True
                )
                client_thread.start()
            except Exception as e:
                if self.running:
                    logger.error(f"Accept error: {e}")
                break

    def handle_client(self, client_socket, addr):
        logger.info(f"New client connected: {addr}")
        while self.running:
            try:
                size_data = client_socket.recv(4)
                if len(size_data) != 4:
                    break
                size = int.from_bytes(size_data, "big")
                if size <= 0:
                    break
                data = b""
                while len(data) < size:
                    packet = client_socket.recv(min(4096, size - len(data)))
                    if not packet:
                        break
                    data += packet
                if len(data) == size:
                    nparr = np.frombuffer(data, np.uint8)
                    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                    if frame is not None:
                        timestamp = time.time()
                        self.clients[addr] = (frame, timestamp)
            except Exception as e:
                logger.error(f"Handle client {addr} error: {e}")
                break
        client_socket.close()
        if addr in self.clients:
            del self.clients[addr]
        logger.info(f"Client disconnected: {addr}")

    def stop(self):
        self.running = False
        if hasattr(self, "server_socket"):
            self.server_socket.close()
        logger.info("VideoServer stopped")


class VideoPanel(QWidget):
    """视频显示面板：选择源、显示帧、标注"""

    def __init__(self, video_server, panel_id, parent=None):
        super().__init__(parent)
        self.video_server = video_server
        self.panel_id = panel_id
        self.annotate_enabled = False
        if USE_YOLO:
            model_path = os.path.abspath(
                os.path.join(os.path.dirname(__file__), "../../datas/yolov11.pt")
            )
            self.model = YOLO(model_path)
        layout = QVBoxLayout(self)
        layout.addWidget(QLabel(f"视频源 {panel_id}:"))
        source_layout = QHBoxLayout()
        self.combo = QComboBox()
        self.combo.addItem("无")
        source_layout.addWidget(self.combo)
        self.annotate_checkbox = QCheckBox("标注")
        self.annotate_checkbox.stateChanged.connect(self.on_annotate_changed)
        source_layout.addWidget(self.annotate_checkbox)
        if not USE_YOLO:
            self.annotate_checkbox.setEnabled(False)
            self.annotate_checkbox.setText("标注(未启用)")
        layout.addLayout(source_layout)
        self.video_label = QLabel()
        self.video_label.setMinimumSize(400, 200)
        self.video_label.setAlignment(Qt.AlignCenter)
        self.video_label.setStyleSheet(
            "border: 1px solid gray; background-color: black;"
        )
        layout.addWidget(self.video_label, 1)
        self.selected_addr = None
        self.combo.currentTextChanged.connect(self.on_selection_changed)
        self.update_timer = QTimer(self)
        self.update_timer.timeout.connect(self.update_combo_and_display)
        self.update_timer.start(100)

    def on_annotate_changed(self, state):
        self.annotate_enabled = state == Qt.Checked

    def set_source(self, addr):
        ip, port = addr
        text = f"{ip}:{port}"
        self.combo.setCurrentText(text)

    def update_combo_and_display(self):
        current_text = self.combo.currentText()
        self.combo.clear()
        self.combo.addItem("无")
        for addr in sorted(self.video_server.clients.keys()):
            ip, port = addr
            self.combo.addItem(f"{ip}:{port}")
        if current_text != "无" and current_text in [
            self.combo.itemText(i) for i in range(1, self.combo.count())
        ]:
            self.combo.setCurrentText(current_text)
        self.update_display()

    def on_selection_changed(self, text):
        if text == "无":
            self.selected_addr = None
            self.video_label.clear()
        else:
            parts = text.split(":")
            if len(parts) == 2:
                self.selected_addr = (parts[0], int(parts[1]))
        self.update_display()

    def update_display(self):
        if self.selected_addr and self.selected_addr in self.video_server.clients:
            frame, ts = self.video_server.clients[self.selected_addr]
            if frame is not None:
                if USE_YOLO and self.annotate_enabled:
                    results = self.model.predict(frame, verbose=False)
                    annotated_frame = results[0].plot(labels=False)
                    frame = annotated_frame
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                height, width, channel = frame_rgb.shape
                bytes_per_line = 3 * width
                q_image = QImage(
                    frame_rgb.data, width, height, bytes_per_line, QImage.Format_RGB888
                )
                pixmap = QPixmap.fromImage(q_image)
                pixmap = pixmap.scaled(
                    self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation
                )
                painter = QPainter(pixmap)
                painter.setPen(QColor(255, 255, 255))
                painter.setFont(QFont("Arial", 10))
                time_str = time.strftime("%H:%M:%S", time.localtime(ts))
                painter.drawText(10, pixmap.height() - 10, f"Updated: {time_str}")
                painter.end()
                self.video_label.setPixmap(pixmap)
        else:
            self.video_label.clear()

if __name__ == "__main__":
    print("启动视频服务器测试...")
    import sys
    from PyQt5.QtWidgets import QApplication
    
    print("UI初始化...")
    # 创建Qt应用
    app = QApplication(sys.argv)
    
    # 创建视频服务器
    video_server = VideoServer()
    
    # 创建主窗口用于显示视频面板
    main_window = QWidget()
    main_layout = QVBoxLayout(main_window)
    
    # 创建两个视频面板实例
    video_panel1 = VideoPanel(video_server, 1)
    video_panel2 = VideoPanel(video_server, 2)
    
    main_layout.addWidget(video_panel1)
    main_layout.addWidget(video_panel2)
    
    main_window.setWindowTitle("视频服务器测试")
    main_window.resize(900, 600)
    main_window.show()
    
    
    print("启动主循环...")
    # 运行应用
    try:
        sys.exit(app.exec_())
    finally:
        # 停止视频服务器
        video_server.stop()